package club.babyf.nativeai.functioncalling.registry;

import club.babyf.nativeai.entity.CustomHashMap;
import club.babyf.nativeai.functioncalling.annotation.Parameters;
import club.babyf.nativeai.functioncalling.annotation.Properties;
import club.babyf.nativeai.functioncalling.annotation.RegisteredFunction;
import club.babyf.nativeai.functioncalling.functions.Function;
import lombok.RequiredArgsConstructor;
import org.reflections.Reflections;
import org.springframework.context.ApplicationContext;
import org.springframework.stereotype.Component;

import java.lang.reflect.Field;
import java.lang.reflect.Method;
import java.util.*;

/**
 * FunctionRegistry.
 *
 * @author zhanglei.
 * @date 2025/2/21 10:46.
 * @description 函数注册器.
 */
@Component
@RequiredArgsConstructor
public class FunctionRegistry {

    private final ApplicationContext applicationContext;

    private final List<CustomHashMap<String, Object>> functionParamList = new ArrayList<>();

    private final Map<String, Function<?, ?>> functions = new HashMap<>(16);

    public void scanAndRegister(String basePackage) throws Exception {
        Reflections reflections = new Reflections(basePackage);
        Set<Class<?>> clazzSet = reflections.getTypesAnnotatedWith(RegisteredFunction.class);
        for (Class<?> clazz : clazzSet) {
            Object newInstance = applicationContext.getBean(clazz);
            RegisteredFunction annotation = clazz.getAnnotation(RegisteredFunction.class);
            String name = annotation.name();
            String description = annotation.description();
            CustomHashMap<String, Object> functionParam = getFunctionParam(clazz, newInstance, name, description);
            functionParamList.add(functionParam);
            functions.put(name, (Function<?, ?>) newInstance);
        }
    }

    public Function<?, ?> getFunction(String name) {
        return functions.get(name);
    }

    public boolean hasFunction(String name) {
        return functions.containsKey(name);
    }

    public List<CustomHashMap<String, Object>> getFunctionParamList() {
        return functionParamList;
    }

    private CustomHashMap<String, Object> getFunctionParam(Class<?> clazz, Object newInstance, String name, String description) throws Exception {
        // 构建函数参数
        CustomHashMap<String, Object> functionParam = new CustomHashMap<>();
        functionParam.put("type", "function");
        // 构建函数
        CustomHashMap<String, Object> function = new CustomHashMap<>();
        function.put("name", name);
        function.put("description", description);
        // 获取方法的Method对象
        Method method = clazz.getMethod("getParamType");
        // 获取方法的返回值
        Class<?> parametersClassType = (Class<?>) method.invoke(newInstance);
        // 获取类型绑定的参数注解
        Parameters parametersAnnotation = parametersClassType.getAnnotation(Parameters.class);
        // 获取参数类型
        String parametersType = parametersAnnotation.type();
        // 构建参数
        CustomHashMap<String, Object> parameters = new CustomHashMap<>();
        parameters.put("type", parametersType);
        // 构建属性实体
        CustomHashMap<String, Object> properties = new CustomHashMap<>();
        // 构建必填属性列表
        List<String> requiredList = new ArrayList<>();
        // 获取参数内的全部属性
        Field[] declaredFields = parametersClassType.getDeclaredFields();
        // 遍历属性
        for (Field declaredField : declaredFields) {
            // 获取属性名称
            String fieldName = declaredField.getName();
            // 获取属性绑定的属性注解
            Properties propertiesAnnotation = declaredField.getAnnotation(Properties.class);
            String propertiesType = propertiesAnnotation.type();
            String propertiesDescription = propertiesAnnotation.description();
            // 构建属性
            CustomHashMap<String, Object> property = new CustomHashMap<>();
            property.put("type", propertiesType);
            property.put("description", propertiesDescription);
            // 将属性放入属性列表
            properties.put(fieldName, property);
            // 判断属性是否必填
            boolean propertiesRequired = propertiesAnnotation.required();
            if (propertiesRequired) {
                requiredList.add(fieldName);
            }
        }
        parameters.put("properties", properties);
        parameters.put("required", requiredList);
        function.put("parameters", parameters);
        functionParam.put("function", function);
        return functionParam;
    }

}
